from pymongo import UpdateOne
from zc_core.client.mongo_client import Mongo
import pandas as pd

mongo = Mongo()


def fix_data_spu():
    data = pd.read_excel('sku_20211014.xls')
    plat_sku_id = list(data['商品编码'].values)
    _id = list(data['_id'].values)
    plat_dict = {}
    for index, data in enumerate(_id):
        plat_dict[str(data)] = str(plat_sku_id[index])
    update_bulk = list()
    src_list = mongo.list('sku_20211016', fields={'productId': 1})
    for item in src_list:
        product_id = item.get('productId')
        _id = item.get('_id')
        if product_id:
            update_bulk.append(UpdateOne({'_id': _id}, {'$set': {
                'platSkuId': plat_dict.get(product_id)
            }}, upsert=False))

    print(len(update_bulk))
    mongo.bulk_write('sku_20211016', update_bulk)
    mongo.close()
    print('任务完成~')


def fix_supplier_name():
    data_list = mongo.distinct('data_20211017', 'supplierId')
    supplier_dict = {}
    update_bulk = list()
    for supplier_id in data_list:
        record = mongo.get('data_20211017',
                           query={'supplierName': {'$exists': True}, 'supplierId': supplier_id})
        supplier_dict[record.get('supplierId')] = record.get('supplierName')
    print(supplier_dict)
    src_list = mongo.list('data_20211017', fields={'supplierId': 1, 'platSkuId': 1})
    for item in src_list:
        supplier_id = item.get('supplierId')
        plat_sku_id = item.get('platSkuId')
        if supplier_id:
            update_bulk.append(UpdateOne({'platSkuId': plat_sku_id}, {'$set': {
                'platSkuId': plat_sku_id,
                'supplierId': supplier_id,
                'supplierName': supplier_dict.get(supplier_id)
            }}, upsert=False))
    print(len(update_bulk))
    mongo.bulk_write('data_20211017', update_bulk)
    print('任务完成')


if __name__ == '__main__':
    # fix_data_spu()
    fix_supplier_name()
